#include "Solver.hpp"
#include <string>
#include <fstream>

//smoothing function using Gauss-Seidel method
void Solver::smoothGS(Array &x, const Array &f, size_t iter)
{
    //assert(x.xsize() == f.xsize()); 

    size_t width = x.xsize();
    real h = x.hSize();

    real omg = 1.0;

    for ( size_t iter_ = 1 ; iter_<=iter ; ++iter_)
    {
        for (size_t i = 1; i <width-1; ++i)
        {
            for (size_t j = 1; j <width-1; ++j)
            {
                x(i,j) = (1.-omg)*x(i,j) + omg*h*h*0.25*( (x(i+1,j)+x(i-1,j)+x(i,j+1)+x(i,j-1))/(h*h) -f(i,j) );
            }
        }
    }

}

//smoothing function using red-black Gauss-Seidel
void Solver::smoothGSRB(Array &x, const Array &f, size_t iter)
{
    //assert(x.xSize() == f.xsize());

    size_t width = x.xsize();
    size_t j;
    real h = x.hSize();
    real omg = 1.0;


    for ( size_t iter_ = 1 ; iter_<=iter ; ++iter_)
    {
        //updating "red" cells
        for (size_t i = 1; i < width-1; ++i)
        {
            (i%2 == 1)?j=1:j=2;
            for (; j < width-1; j+=2)
            {
                x(i,j) = (1-omg)*x(i,j) + omg*h*h*0.25*( (x(i+1,j)+x(i-1,j)+x(i,j+1)+x(i,j-1))/(h*h) -f(i,j) );
            }
        }
        //updating "black" cells
        for (size_t i = 1; i < width-1; ++i)
        {
            (i%2 == 1)?j = 2:j=1;
            for (; j < width-1; j+=2)
            {
                x(i,j) = (1-omg)*x(i,j) + omg*h*h*0.25*( (x(i+1,j)+x(i-1,j)+x(i,j+1)+x(i,j-1))/(h*h) -f(i,j) );
            }
        }
    }
}

//restriction function
void Solver::restrict(Array &r, Array &f)
{
    const size_t size ( f. xsize () );

    for ( size_t i=1; i<size-1; ++i )
    {
        for ( size_t j=1; j<size-1; ++j )
        {
            f(i,j) = 1./4.  * r(i*2    , j*2  )  +
                     1./8.  * r(i *2+1 , j*2  )  +
                     1./8.  * r(i*2 -1 , j*2  )  +
                     1./8.  * r(i*2    , j *2 -1)+
                     1./8.  * r(i*2    , j *2+1) +
                     1./16. * r(i *2+1 , j *2 -1)+
                     1./16. * r(i *2+1 , j *2+1) +
                     1./16. * r(i*2 -1 , j *2 -1)+
                     1./16. * r(i*2 -1 , j *2+1);
        }
    }
}

/*
   function for calculating the residual;
   Ax-f =r (the residual is not elevated to the second power)
*/
void Solver::calcResidual(const Array &x, const Array &f, Array &r)
{
    const size_t size ( f. xsize () );

    for(size_t i=1; i<size-1; ++i)
    {
        for(size_t j=1; j<size-1; ++j)
        {
            r(i,j)=-(x(i+1,j)-4.0*x(i,j)+x(i-1,j)+x(i,j+1)+x(i,j-1))/(x.hSize()*x.hSize())+f(i,j);
        }
    }
}

void Solver::prolongateAndCorrect(const Array &e, Array &x)
{
    Array y(2);
    y=x;
    y.fill(0.0);
    const size_t size(e. xsize ());

    for ( size_t i =1; i<size -1; ++i )
    {
        for ( size_t j =1; j<size -1; ++j )
        {
            y(2*i,2*j) += e(i,j);
            y(2*i+1,2*j) += (e(i,j))*0.5;
            y(2*i-1,2*j) += (e(i,j))*0.5;
            y(2*i,2*j+1) += (e(i,j))*0.5;
            y(2*i,2*j-1) += (e(i,j))*0.5;
            y(2*i+1,2*j+1) += (e(i,j))*0.25;
            y(2*i-1,2*j+1) += (e(i,j))*0.25;
            y(2*i+1,2*j-1) += (e(i,j))*0.25;
            y(2*i-1,2*j-1) += (e(i,j))*0.25;
			
			            
						x(2*i,2*j) += e(i,j)*1./8.;
            x(2*i+1,2*j) += (e(i,j)+e(i+1,j))*1./8.;
            //x(2*i-1,2*j) += (e(i,j)+e(i-1,j))*1/4;
            x(2*i,2*j+1) += (e(i,j)+e(i,j+1))*1./8.;
            //x(2*i,2*j-1) += (e(i,j)+e(i,j-1))*1/4;
            x(2*i+1,2*j+1) += (e(i,j)+e(i,j+1)+e(i+1,j)+e(i+1,j+1))*1./16.;
 //           x(2*i-1,2*j+1) += (e(i,j)+e(i,j+1)+e(i-1,j)+e(i-1,j+1))*1/8;
   //         x(2*i+1,2*j-1) += (e(i,j)+e(i,j-1)+e(i+1,j)+e(i+1,j-1))*1/8;
     //       x(2*i-1,2*j-1) += (e(i,j)+e(i,j-1)+e(i-1,j)+e(i-1,j-1))*1/8;
			
			
						         */   
	x(2*i,2*j) += e(i,j);
            x(2*i+1,2*j) += (e(i,j))*0.5;
            x(2*i-1,2*j) += (e(i,j))*0.5;
            x(2*i,2*j+1) += (e(i,j))*0.5;
            x(2*i,2*j-1) += (e(i,j))*0.5;
            x(2*i+1,2*j+1) += (e(i,j))*0.25;
            x(2*i-1,2*j+1) += (e(i,j))*0.25;
            x(2*i+1,2*j-1) += (e(i,j))*0.25;
            x(2*i-1,2*j-1) += (e(i,j))*0.25;
        }
    }

    for ( size_t i =0; i<x.xsize()-1; ++i )
    {
        for ( size_t j =0; j<x.xsize()-1; ++j )
        {
            x(i,j) +=y(i,j);
        }
    }

}

//multigrid solver function
void Solver::mgSolve(Array& x, const Array& f, size_t cycles, size_t pre_iter, size_t post_iter)
{
    size_t finest_l=levels_-1;
    real res1,res2;

    std::string filename1="error.txt";
    std::ofstream file1(filename1);

    file1<<"Cycle\t\tL2_Res\t\tConvergence_Ratio\t\tError_L2\n";
    std::cout<<"Cycle\t\tL2_Res\t\tConvergence_Ratio\t\tError_L2\n";
    for(size_t c=0; c<cycles; ++c)
    {
        //Finest grid pre-smoothing
        smoothGSRB(x,f, pre_iter);
        //Finest grid residual calculation
        calcResidual(x,f, r_[finest_l]);
        res1=r_[finest_l].l2Norm();
        //Restrict the residual at finest level
        restrict(r_[finest_l], f_[finest_l-1]);
        //Initialize the next grid
        x_[finest_l-1].fill(0.0);

        //Intermediate grids
        for(size_t l=finest_l-1; l>1; --l)
        {
            //pre-smoothing
            smoothGSRB(x_[l],f_[l], pre_iter);

            //Residual calculation
            calcResidual(x_[l],f_[l],r_[l]);
            //Restricting the residual at coarser grids
            restrict(r_[l],f_[l-1]);
            //Initialize the next grid
            x_[l-1].fill(0.0);

        }
        
        //Solving coarsest grid
        smoothGS(x_[1],f_[1],1);
        //Intermediate levels
        for(size_t l=2; l<finest_l; ++l)
        {
            //Prolongate and correct
            prolongateAndCorrect(x_[l-1],x_[l]);
            //Post-smoothing
            smoothGSRB(x_[l],f_[l],post_iter);
        }

        //Prolongate and correct finest grid
        prolongateAndCorrect(x_[finest_l-1],x);
        //Smooth finest grid
        smoothGSRB(x,f,post_iter);
        //Calculate residual
        calcResidual(x,f,r_[finest_l]);
        res2=r_[finest_l].l2Norm();
        
        file1<<c+1<<"\t\t"<<r_[finest_l].l2Norm()<< "\t\t"<<res2/res1<<"\t\t" <<x_[finest_l-1].l2Norm()<<"\n";
        std::cout<<c+1<<"\t\t"<<r_[finest_l].l2Norm()<< "\t\t"<<res2/res1<<"\t\t" <<x_[finest_l-1].l2Norm()<<"\n";
    }
}


void Solver::prolongate(const Array &e, Array &x)
{
	Array y(2);
	y=x;
	const size_t size ( e. xsize () );
    for ( size_t i =1; i<size -1; ++i ) 
    {
        for ( size_t j =1; j<size -1; ++j )
        {
  						y(2*i,2*j) = e(i,j);
            y(2*i+1,2*j) = (e(i,j)+e(i+1,j))*0.5;
            y(2*i-1,2*j) = (e(i,j)+e(i-1,j))*0.5;
            y(2*i,2*j+1) = (e(i,j)+e(i,j+1))*0.5;
            y(2*i,2*j-1) = (e(i,j)+e(i,j-1))*0.5;
            y(2*i+1,2*j+1) = (e(i,j)+e(i,j+1)+e(i+1,j)+e(i+1,j+1))*0.25;
           y(2*i-1,2*j+1) = (e(i,j)+e(i,j+1)+e(i-1,j)+e(i-1,j+1))*0.25;
            y(2*i+1,2*j-1) = (e(i,j)+e(i,j-1)+e(i+1,j)+e(i+1,j-1))*0.25;
            y(2*i-1,2*j-1) = (e(i,j)+e(i,j-1)+e(i-1,j)+e(i-1,j-1))*0.25;
		
        }
    }
	//y.print(std::cout);
	for ( size_t i =1; i<x.xsize()-1; ++i ) 
    {
        for ( size_t j =1; j<x.xsize()-1; ++j )
        {
		x(i,j) +=y(i,j);
		}
	}	
	
}




